#!/usr/bin/python
# -*- encoding:utf-8 -*-

"""
    使用cryptography库提供的AES加密
"""
from cryptography.fernet import Fernet
import base64
import os
from sftp_account_manage.com.exception.NullException import NPException


# 加密解密的工具类
class Encrypt(object):

    def __init__(self):
        self.key = Fernet.generate_key()

    def set_key(self, key: bytes):
        # the length of this key msut be a multiple of three, if not, error will occur ---- LiuShuai Noted on 30 Dec
        self.key = key

    def get_encrypt_str(self, text: bytes):
        # try:
        if self.key == '':
            raise NPException("[ERROR]: 密匙未初始化")
        token = Fernet(self.key).encrypt(text)
        return self.bytes_to_str(token)
        # except NPException as e:
        #     print(e.info)

    def get_decrypt_str(self, token: bytes):
        if self.key == '':
            raise NPException("[ERROR]: 密匙未初始化")
        try:
            text = Fernet(self.key).decrypt(token)
        except Exception as e:
            return ''
        return self.bytes_to_str(text)


    @staticmethod
    # 生成一个bytes类型的key,用于做对称加密的密匙
    def generate_key():
        return base64.urlsafe_b64encode(os.urandom(32))
        # return Fernet.generate_key()

    @staticmethod
    # 将string转换成bytes
    def str_to_bytes(this_text: str):
        return str.encode(this_text)

    @staticmethod
    # 将bytes转换成string
    def bytes_to_str(this_byte: bytes):
        return bytes.decode(this_byte)
